import sys
import argparse
import json
import pika

from mlxtend.frequent_patterns import fpgrowth, association_rules
import pandas as pd
from collections import Counter

# 解析命令行参数
def parse_args():
    parser = argparse.ArgumentParser(description="FP-Growth Algorithm for Association Rules")
    parser.add_argument('--rabbitmq-host', type=str, default='localhost', help="RabbitMQ host")
    parser.add_argument('--input-queue', type=str, default='FP_Growth_Input', help="Input queue name")
    return parser.parse_args()

# 过滤低频项目
def filter_low_frequency_items(data, min_frequency=5):
    counter = Counter(item for transaction in data for item in transaction)
    filtered_data = [
        [item for item in transaction if counter[item] >= min_frequency]
        for transaction in data
    ]
    return filtered_data

# 将数据转换为适合 FP-Growth 算法的格式
def convert_to_one_hot(data):
    all_items = set(item for transaction in data for item in transaction)
    encoded_data = []
    for transaction in data:
        row = {item: (item in transaction) for item in all_items}
        encoded_data.append(row)
    return pd.DataFrame(encoded_data)

# 统计符合规则的输入出现次数
def count_rule_occurrences(rules, raw_data):
    rule_counts = {}
    for _, rule in rules.iterrows():
        antecedents = set(rule['antecedents'])
        consequents = set(rule['consequents'])

        # 初始化计数
        rule_key = (frozenset(antecedents), frozenset(consequents))
        rule_counts[rule_key] = 0

        # 遍历原始数据，统计符合条件的事务
        for transaction in raw_data:
            transaction_set = set(transaction)
            if antecedents.issubset(transaction_set) and consequents.issubset(transaction_set):
                rule_counts[rule_key] += 1

    return rule_counts

# 处理消息
def process_message(channel, method, properties, body):
    try:
        # 打印原始消息
        # print(f"Received raw message: {body}")

        # 解码消息
        try:
            message = json.loads(body.decode('utf-8'))
        except (json.JSONDecodeError, UnicodeDecodeError) as e:
            print(f"Error decoding message: {e}", file=sys.stderr)
            channel.basic_ack(delivery_tag=method.delivery_tag)
            return

        # 检查消息格式
        if not isinstance(message, dict) or "data" not in message:
            print("Error: Invalid message format", file=sys.stderr)
            channel.basic_ack(delivery_tag=method.delivery_tag)
            return

        data = message["data"]
        if not isinstance(data, dict) or "support" not in data or "confidence" not in data or "data" not in data:
            print("Error: Missing required fields in message", file=sys.stderr)
            channel.basic_ack(delivery_tag=method.delivery_tag)
            return

        # 提取消息中的数据和参数
        support = data["support"]
        confidence = data["confidence"]
        min_frequency = data.get("min_frequency", 5)
        raw_data = [line.split(',') for line in data["data"].strip().split('\n')]

        # 过滤低频项目
        raw_data = filter_low_frequency_items(raw_data, min_frequency)

        # 数据预处理：转换为 one-hot 编码
        df = convert_to_one_hot(raw_data)

        # 使用 FP-Growth 算法找出频繁项集
        frequent_itemsets = fpgrowth(df, min_support=support, use_colnames=True)

        # 检查频繁项集是否为空
        if frequent_itemsets.empty:
            # 如果频繁项集为空，则返回空结果
            output_message = json.dumps({"pattern": "FP_Growth_Result", "data": []})
            channel.basic_publish(
                exchange='',
                routing_key=properties.reply_to,  # 发送到客户端指定的回调队列
                body=output_message,
                properties=pika.BasicProperties(
                    correlation_id=properties.correlation_id  # 匹配关联 ID
                )
            )
            # 确认消息已被处理
            channel.basic_ack(delivery_tag=method.delivery_tag)
            return

        # 生成关联规则
        rules = association_rules(frequent_itemsets, metric="confidence", min_threshold=confidence)

        # 按提升度排序
        rules = rules.sort_values(by='lift', ascending=False)

        # 统计符合规则的输入出现次数
        rule_counts = count_rule_occurrences(rules, raw_data)

        # 格式化输出结果
        result = []
        for _, rule in rules.iterrows():
            antecedents = set(rule['antecedents'])
            consequents = set(rule['consequents'])
            confidence = rule['confidence']
            lift = rule['lift']

            # 获取该规则的出现次数
            rule_key = (frozenset(antecedents), frozenset(consequents))
            occurrence_count = rule_counts.get(rule_key, 0)

            # 构建规则字典
            rule_dict = {
                "antecedents": list(antecedents),
                "consequents": list(consequents),
                "confidence": round(confidence, 2),
                "lift": round(lift, 2),
                "occurrence_count": occurrence_count
            }
            result.append(rule_dict)

        # 将结果封装为 JSON 并发送到回调队列
        output_message = json.dumps({"pattern": "FP_Growth_Result", "data": result})
        channel.basic_publish(
            exchange='',
            routing_key=properties.reply_to,  # 发送到客户端指定的回调队列
            body=output_message,
            properties=pika.BasicProperties(
                correlation_id=properties.correlation_id  # 匹配关联 ID
            )
        )

        # 确认消息已被处理
        channel.basic_ack(delivery_tag=method.delivery_tag)

    except Exception as e:
        print(f"Error processing message: {str(e)}", file=sys.stderr)
        # 如果处理失败，拒绝消息并重新入队
        channel.basic_nack(delivery_tag=method.delivery_tag, requeue=True)

# 主函数
def main():
    try:
        # 解析命令行参数
        args = parse_args()
        rabbitmq_host = args.rabbitmq_host
        input_queue = args.input_queue

        # 连接到 RabbitMQ
        connection = pika.BlockingConnection(pika.ConnectionParameters(host=rabbitmq_host))
        channel = connection.channel()

        # 声明输入队列（持久化）
        channel.queue_declare(queue=input_queue, durable=True)

        # 设置消费者
        channel.basic_consume(
            queue=input_queue,
            on_message_callback=process_message
        )

        print(f"Waiting for messages in queue '{input_queue}'...")
        channel.start_consuming()

    except KeyboardInterrupt:
        print("Interrupted by user. Exiting...")
        sys.exit(0)
    except Exception as e:
        print(f"Error: {str(e)}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()